package com.study.spark.scala.ml.cluster

import org.apache.spark.ml.clustering.LDA
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.SparkSession

import scala.util.Random

/**
  * 聚类：LDA算法，基于概率统计
  * LDA及文档主题生成模型，属于无监督学习
  * 用于将文档分成K个主题
  *
  * @author stephen
  * @create 2019-04-07 20:25
  * @since 1.0.0
  */
object LDADemo {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("LDA Demo")
      .master("local[2]")
      .getOrCreate()
    val file = spark.read
      .format("csv")
      .option("sep", ",")
      .load("/Users/stephen/Documents/03code/java-demo/bigdata-study/study-spark/src/main/resource/classification/iris.data")

    val random = new Random()
    import spark.implicits._
    val data = file.map(row => {
      val label = row.getString(4) match {
        case "Iris-setosa" => 0
        case "Iris-versicolor" => 1
        case "Iris-virginica" => 2
      }
      (
        row.getString(0).toDouble,
        row.getString(1).toDouble,
        row.getString(2).toDouble,
        row.getString(3).toDouble,
        label,
        random.nextDouble()
      )
    }).toDF("_c0", "_c1", "_c2", "_c3", "label", "random").sort("random")

    val assembler = new VectorAssembler()
      .setInputCols(Array("_c0", "_c1", "_c2", "_c3"))
      .setOutputCol("features")
    val dataset = assembler.transform(data)
    val Array(train, test) = dataset.randomSplit(Array(0.8, 0.2))

    // 基于概率统计的算法的迭代次数一般比较大
    val lda = new LDA().setFeaturesCol("features").setK(3).setMaxIter(30)
    val model = lda.fit(train)
    val prediction = model.transform(train)
    // 最大似然估计，越大越好
    val ll = model.logLikelihood(train)
    // 疑惑度，越小越好
    val lp = model.logPerplexity(train)

    val topics = model.describeTopics(3)
    prediction.select("label","topicDistribution").show(false)
    printf("The topics described by their top-weighted terms:")
    topics.show(false)

    println(s"The lower bound on the log likelihood of the entire corpus:$ll")
    println(s"The upper bound on perplexity:$lp")
  }
}
